{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from sentence_transformers.evaluation import (\n",
    "    InformationRetrievalEvaluator,\n",
    "    SequentialEvaluator,\n",
    ")\n",
    "from sentence_transformers.util import cos_sim\n",
    "from datasets import load_dataset, concatenate_datasets\n",
    "model_id = \"sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2\"  # Hugging Face model ID\n",
    "matryoshka_dimensions = [384, 256, 128, 64] # Important: large to small\n",
    "model = SentenceTransformer(\n",
    "    model_id, device=\"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    ")\n",
    "test_dataset = load_dataset(\"json\", data_files=\"all-data/test_en.json\", split=\"train\")\n",
    "train_dataset = load_dataset(\"json\", data_files=\"all-data/train_en.json\", split=\"train\")\n",
    "#更改属性名\n",
    "test_dataset = test_dataset.rename_column(\"text\", \"anchor\")\n",
    "test_dataset = test_dataset.rename_column(\"label\", \"positive\")\n",
    "train_dataset = train_dataset.rename_column(\"text\", \"anchor\")\n",
    "train_dataset = train_dataset.rename_column(\"label\", \"positive\")\n",
    "corpus_dataset = concatenate_datasets([train_dataset, test_dataset])\n",
    "corpus = dict(\n",
    "    zip([range(1,len(corpus_dataset[\"positive\"])+1)], corpus_dataset[\"positive\"])\n",
    ")  # Our corpus (cid => document)\n",
    "queries = dict(\n",
    "    zip([range(1,len(corpus_dataset[\"anchor\"])+1)], corpus_dataset[\"anchor\"])\n",
    ")  # Our queries (qid => question)\n",
    "relevant_docs = {}  # Query ID to relevant documents (qid => set([relevant_cids])\n",
    "for q_id in queries:\n",
    "    relevant_docs[q_id] = [q_id]\n",
    "\n",
    "\n",
    "matryoshka_evaluators = []\n",
    "# Iterate over the different dimensions\n",
    "for dim in matryoshka_dimensions:\n",
    "    ir_evaluator = InformationRetrievalEvaluator(\n",
    "        queries=queries,\n",
    "        corpus=corpus,\n",
    "        relevant_docs=relevant_docs,\n",
    "        name=f\"dim_{dim}\",\n",
    "        truncate_dim=dim,  # Truncate the embeddings to a certain dimension\n",
    "        score_functions={\"cosine\": cos_sim},\n",
    "    )\n",
    "    matryoshka_evaluators.append(ir_evaluator)\n",
    "\n",
    "# Create a sequential evaluator\n",
    "evaluator = SequentialEvaluator(matryoshka_evaluators)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dim_384_cosine_ndcg@10: 1.0\n",
      "dim_256_cosine_ndcg@10: 1.0\n",
      "dim_128_cosine_ndcg@10: 1.0\n",
      "dim_64_cosine_ndcg@10: 1.0\n"
     ]
    }
   ],
   "source": [
    "# Evaluate the model\n",
    "results = evaluator(model)\n",
    "\n",
    "# # COMMENT IN for full results\n",
    "# print(results)\n",
    "\n",
    "# Print the main score\n",
    "for dim in matryoshka_dimensions:\n",
    "    key = f\"dim_{dim}_cosine_ndcg@10\"\n",
    "    print\n",
    "    print(f\"{key}: {results[key]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformerModelCardData, SentenceTransformer\n",
    "\n",
    "# Hugging Face model ID: https://huggingface.co/BAAI/bge-base-en-v1.5\n",
    "model_id = \"sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2\"\n",
    "\n",
    "# load model with SDPA for using Flash Attention 2\n",
    "model = SentenceTransformer(\n",
    "    model_id,\n",
    "    model_kwargs={\"attn_implementation\": \"sdpa\"},\n",
    "    model_card_data=SentenceTransformerModelCardData(\n",
    "        language=\"multilingual\",#\"en\",\n",
    "        license=\"apache-2.0\",\n",
    "        model_name=\"distiluse_basev2 Financial GND\",\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers.losses import MatryoshkaLoss, MultipleNegativesRankingLoss\n",
    "\n",
    "matryoshka_dimensions = [384, 256, 128, 64]  # Important: large to small\n",
    "inner_train_loss = MultipleNegativesRankingLoss(model)\n",
    "train_loss = MatryoshkaLoss(\n",
    "    model, inner_train_loss, matryoshka_dims=matryoshka_dimensions\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformerTrainingArguments\n",
    "from sentence_transformers.training_args import BatchSamplers\n",
    "\n",
    "# load train dataset again\n",
    "train_dataset = load_dataset(\"json\", data_files=\"all-data/train_en.json\", split=\"train\")\n",
    "#更改属性名\n",
    "train_dataset = train_dataset.rename_column(\"text\", \"anchor\")\n",
    "train_dataset = train_dataset.rename_column(\"label\", \"positive\")\n",
    "# define training arguments\n",
    "args = SentenceTransformerTrainingArguments(\n",
    "    output_dir=\"multilingual-MiniLM-L12-v2\", # output directory and hugging face model ID\n",
    "    num_train_epochs=4,                         # number of epochs\n",
    "    per_device_train_batch_size=32,             # train batch size\n",
    "    gradient_accumulation_steps=16,             # for a global batch size of 512\n",
    "    per_device_eval_batch_size=16,              # evaluation batch size\n",
    "    warmup_ratio=0.1,                           # warmup ratio\n",
    "    learning_rate=2e-5,                         # learning rate, 2e-5 is a good value\n",
    "    lr_scheduler_type=\"cosine\",                 # use constant learning rate scheduler\n",
    "    optim=\"adamw_torch_fused\",                  # use fused adamw optimizer\n",
    "    tf32=True,                                  # use tf32 precision\n",
    "    bf16=True,                                  # use bf16 precision\n",
    "    batch_sampler=BatchSamplers.NO_DUPLICATES,  # MultipleNegativesRankingLoss benefits from no duplicate samples in a batch\n",
    "    eval_strategy=\"epoch\",                      # evaluate after each epoch\n",
    "    save_strategy=\"epoch\",                      # save after each epoch\n",
    "    logging_steps=10,                           # log every 10 steps\n",
    "    save_total_limit=3,                         # save only the last 3 models\n",
    "    load_best_model_at_end=True,                # load the best model when training ends\n",
    "    metric_for_best_model=\"eval_dim_128_cosine_ndcg@10\",  # Optimizing for the best ndcg@10 score for the 128 dimension\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformerTrainer\n",
    "\n",
    "trainer = SentenceTransformerTrainer(\n",
    "    model=model, # bg-base-en-v1\n",
    "    args=args,  # training arguments\n",
    "    train_dataset=train_dataset.select_columns(\n",
    "        [\"positive\", \"anchor\"]\n",
    "    ),  # training dataset\n",
    "    loss=train_loss,\n",
    "    evaluator=evaluator,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "db26e2a344304acd8ec5a8078a2cfa07",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/284 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Column 'anchor' is at index 1, whereas a column with this name is usually expected at index 0. Note that the column order can be important for some losses, e.g. MultipleNegativesRankingLoss will always consider the first column as the anchor and the second as the positive, regardless of the dataset column names. Consider renaming the columns to match the expected order, e.g.:\n",
      "dataset = dataset.select_columns(['anchor', 'positive', 'negative'])\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 5.2664, 'grad_norm': 34.90868377685547, 'learning_rate': 6.896551724137932e-06, 'epoch': 0.14}\n",
      "{'loss': 2.8554, 'grad_norm': 12.509549140930176, 'learning_rate': 1.3793103448275863e-05, 'epoch': 0.28}\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[32], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m \u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m      3\u001b[0m \u001b[38;5;66;03m# save the best model\u001b[39;00m\n\u001b[0;32m      4\u001b[0m trainer\u001b[38;5;241m.\u001b[39msave_model()\n",
      "File \u001b[1;32m~\\AppData\\Roaming\\Python\\Python310\\site-packages\\transformers\\trainer.py:2052\u001b[0m, in \u001b[0;36mTrainer.train\u001b[1;34m(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)\u001b[0m\n\u001b[0;32m   2050\u001b[0m         hf_hub_utils\u001b[38;5;241m.\u001b[39menable_progress_bars()\n\u001b[0;32m   2051\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 2052\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43minner_training_loop\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m   2053\u001b[0m \u001b[43m        \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   2054\u001b[0m \u001b[43m        \u001b[49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   2055\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtrial\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtrial\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   2056\u001b[0m \u001b[43m        \u001b[49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   2057\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\AppData\\Roaming\\Python\\Python310\\site-packages\\transformers\\trainer.py:2393\u001b[0m, in \u001b[0;36mTrainer._inner_training_loop\u001b[1;34m(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)\u001b[0m\n\u001b[0;32m   2387\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39maccelerator\u001b[38;5;241m.\u001b[39maccumulate(model):\n\u001b[0;32m   2388\u001b[0m     tr_loss_step \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtraining_step(model, inputs)\n\u001b[0;32m   2390\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[0;32m   2391\u001b[0m     args\u001b[38;5;241m.\u001b[39mlogging_nan_inf_filter\n\u001b[0;32m   2392\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m is_torch_xla_available()\n\u001b[1;32m-> 2393\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m (torch\u001b[38;5;241m.\u001b[39misnan(tr_loss_step) \u001b[38;5;129;01mor\u001b[39;00m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43misinf\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtr_loss_step\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[0;32m   2394\u001b[0m ):\n\u001b[0;32m   2395\u001b[0m     \u001b[38;5;66;03m# if loss is nan or inf simply add the average of previous logged losses\u001b[39;00m\n\u001b[0;32m   2396\u001b[0m     tr_loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m tr_loss \u001b[38;5;241m/\u001b[39m (\u001b[38;5;241m1\u001b[39m \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstate\u001b[38;5;241m.\u001b[39mglobal_step \u001b[38;5;241m-\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_globalstep_last_logged)\n\u001b[0;32m   2397\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "trainer.train()\n",
    "\n",
    "# save the best model\n",
    "trainer.save_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "fine_tuned_model = SentenceTransformer(\n",
    "    args.output_dir, device=\"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    ")\n",
    "# Evaluate the model\n",
    "results = evaluator(fine_tuned_model)\n",
    "\n",
    "# # COMMENT IN for full results\n",
    "# print(results)\n",
    "\n",
    "# Print the main score\n",
    "for dim in matryoshka_dimensions:\n",
    "    key = f\"dim_{dim}_cosine_ndcg@10\"\n",
    "    print(f\"{key}: {results[key]}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py311",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
